# interpolation.py
# 插值模块，包括IDW插值和Kriging插值
import os
import numpy as np
from pykrige.ok import OrdinaryKriging
from config import Config
from scipy.spatial import distance_matrix
import pandas as pd

# 球面距离计算
def haversine_vectorized(lon1, lat1, lon2_array, lat2_array):
    """
    计算一个点到多个点的球面距离（单位：km），输入可为numpy数组
    """
    R = 6371.0  # 地球半径，单位km
    lon1, lat1 = np.radians(lon1), np.radians(lat1)
    lon2_array, lat2_array = np.radians(lon2_array), np.radians(lat2_array)

    dlon = lon2_array - lon1
    dlat = lat2_array - lat1
    a = np.sin(dlat / 2.0) ** 2 + np.cos(lat1) * np.cos(lat2_array) * np.sin(dlon / 2.0) ** 2
    c = 2 * np.arcsin(np.sqrt(a))
    return R * c

# IDW插值算法
def idw_interpolation(x, y, z, xi, yi, power=2):
    """
    反距离加权（IDW）插值

    参数:
    x, y : 已知站点的经纬度
    z : 已知站点的降水量
    xi, yi : 待插值站点的经纬度
    power : IDW的幂次，决定了距离对权重的影响

    返回:
    插值结果
    """
    dist = haversine_vectorized(xi, yi, x, y)

    dist[dist == 0] = 1e-10  # 避免除零错误
    weights = 1 / (dist ** power)  # 计算权重
    return np.sum(weights * z) / np.sum(weights)  # 返回加权平均值

# 批量执行IDW插值
def idw_batch(df):
    preds = []
    coords = df[["lon", "lat"]].values
    values = df["pre"].values
    for i in range(len(df)):
        xi, yi = coords[i]
        x = np.delete(coords[:, 0], i)
        y = np.delete(coords[:, 1], i)
        z = np.delete(values, i)

        # 根据插值点动态选取邻近站点
        x_neigh, y_neigh, z_neigh = get_neighbors_within_radius(
            x, y, z, xi, yi,
            min_points=Config.IWD_MIN_POINTS,
            init_radius=Config.IDW_SEARCH_RADIUS,
            step=Config.IDW_SEARCH_STEP,
            max_radius=25
        )

        power = get_idw_power(values[i])
        pred = idw_interpolation(x_neigh, y_neigh, z_neigh, xi, yi, power=power)
        preds.append(pred)
    return preds

# IDW动态搜索半径以获取至少min_points个邻近站点
def get_neighbors_within_radius(x, y, z, xi, yi, min_points=5, init_radius=5, step=2, max_radius=50):
    """
    动态扩大搜索半径，直到找到至少min_points个邻近站点
    参数：
        x, y, z：原始所有站点的经纬度和降水值
        xi, yi：当前待插值站点的经纬度
    返回：
        限定范围内的邻近站点x, y, z
    """
    radius = init_radius
    while radius <= max_radius:
        # 使用向量化haversine计算所有站点与当前点的距离
        dists = haversine_vectorized(xi, yi, x, y)
        mask = dists <= radius
        if np.sum(mask) >= min_points:
            return x[mask], y[mask], z[mask]
        radius += step
    # 若到最大半径仍不足min_points，返回所有站点（此时结果可能不理想，但不为空）
    return x, y, z

# 根据降水量动态调整IDW幂次
def get_idw_power(rain):
    """
    根据降水量动态调整IDW幂次

    参数:
    rain : 降水量

    返回:
    对应的幂次值
    """
    if rain < 10:
        return Config.IDW_POWER_LOW  # 小于10mm时使用较低的幂次
    elif rain < 30:
        return Config.IDW_POWER_MED  # 10-30mm时使用中等的幂次
    else:
        return Config.IDW_POWER_HIGH  # 大于30mm时使用较高的幂次

# Kriging插值算法
def get_closest_neighbors(x, y, z, max_neighbors, xi, yi, df, index, all_neighbor_data):
    """
    获取距离待插值点最近的邻域站点
    
    参数:
    x, y : 已知站点的经纬度
    z : 已知站点的降水量
    max_neighbors : 最大邻域站点数量
    xi, yi : 待插值站点的经纬度
    df : 原始数据 DataFrame，包含站号信息
    index : 当前待插值点在 df 中的索引
    all_neighbor_data : 存储所有邻域站点信息的列表
    
    返回:
    截取后的邻域站点的经纬度和降水量数据
    """
    coords = np.vstack((x, y)).T  # 已知站点的经纬度
    distances = haversine_vectorized(xi, yi, coords[:, 0], coords[:, 1])
    sorted_indices = np.argsort(distances)  # 按距离排序

    # 获取站号
    station_ids = df["Station_Id_C"].drop(index).values[sorted_indices[:max_neighbors]]
    closest_x = x[sorted_indices[:max_neighbors]]  # 选择最近的邻域站点
    closest_y = y[sorted_indices[:max_neighbors]]
    closest_z = z[sorted_indices[:max_neighbors]]
    closest_distances = distances[sorted_indices[:max_neighbors]]

    current_station_id = df.loc[index, "Station_Id_C"]
    for i in range(len(station_ids)):
        all_neighbor_data.append({
            "当前处理站点站号": current_station_id,
            "周边站点站号": station_ids[i],
            "距离(km)": closest_distances[i],
            "降水量": closest_z[i]
        })

    return closest_x, closest_y, closest_z

# 修改 kriging_batch 函数以传递 df 和 index
def kriging_batch(df):
    """
    批量执行Kriging插值
    
    参数:
    df : 数据框，包含站点信息及降水量数据
    
    返回:
    插值结果列表、标准差列表
    """
    preds = []
    stds = []
    coords = df[["lon", "lat"]].values  # 站点经纬度
    values = df["pre"].values  # 站点降水量
    all_neighbor_data = []

    for i in range(len(df)):
        # 每个待估点，排除自身站点
        x = np.delete(coords[:, 0], i)
        y = np.delete(coords[:, 1], i)
        z = np.delete(values, i)
        xi, yi = coords[i]  # 待插值点的经纬度
        
        # 限制Kriging插值的邻域站点数量
        max_neighbors = get_kriging_neighbors(values[i])  # 根据降水量动态调整最大邻域站点数量
        x, y, z = get_closest_neighbors(x, y, z, max_neighbors, xi, yi, df, i, all_neighbor_data)

        try:
            OK = OrdinaryKriging(x, y, z, variogram_model="spherical", verbose=False, enable_plotting=False)  # linear, power, gaussian, spherical, exponential 指数, hole-effect
            z_val, z_std = OK.execute("points", [xi], [yi])
            # 确保估算值非负
            preds.append(max(0, float(z_val[0])))  # 插值结果
            stds.append(float(z_std[0]))  # 标准差
        except:
            preds.append(np.nan)  # 插值失败时返回NaN
            stds.append(np.nan)

    # 将所有邻域站点信息保存到一个 Excel 文件
    neighbor_df = pd.DataFrame(all_neighbor_data)
    excel_file_name = os.path.join("cenresult", f"all_neighbor_info.xlsx")
    neighbor_df.to_excel(excel_file_name, index=False)

    return preds, stds

def get_kriging_neighbors(rain):
    if rain < 10:
        return 20  # 小雨时使用最多邻站，区域较强
    elif rain < 30:
        return 15  # 中雨时使用15~20邻站，区域较弱
    else:
        return 10  # 强降水局地较强，站点控制在10~15更合适